'''
Created on Feb 1, 2012

@author: xzhu
'''

def combine_list(all_lists):
    '''
    combine all the lists
    '''
    # Example:
    #    all_lists = [['-', '0'], ['0', '+']]
    
    # all_lists is not empty
    all_states = []

    if all_lists:
        for i in all_lists[0]:                          # all_lists[0] = ['-', '0']
            all_states.append([i])                      # all_states = [['-'], ['0']]
            
        all_lists.pop(0)
        
        # append each element of the rest lists, generating new states
        for ilist in all_lists:                         # ilist = ['0', '+']
            buff_states = []
            
            for state in all_states:                    # state = ['-']
                for i in ilist:                         # i = '0'
                    new_state = state[:]                # new_state = ['-']
                    new_state.append(i)                 # new_state = ['-', '0'] 
                    buff_states.append(new_state)       # buff_states = [['-', '0']]
                    
            all_states = buff_states[:]
                    
    return all_states

def ave_list(lt):
    if not lt: return 0.0
    else:
        listSum = 0.0
        for i in lt:
            listSum += i
            
        return listSum/len(lt)

import math
import const    
def boltz_ave_list(lt):
    '''
    boltzmann average of a distribution
    '''
    
    if not lt: return 0.0
    
    eRef = lt[0]
    for i in lt:
        if i < eRef: eRef = i
        
    deSum = 0.0
    nuSum = 0.0
    for i in lt:
        deSum += math.exp(-(i - eRef) * const.Kcal2kT)
        nuSum += (i * math.exp(-(i - eRef) * const.Kcal2kT))
        
    return nuSum/deSum    
        
     
    
def std_list(lt):
    
    if (not lt) or len(lt) == 1: return 0.0
    else:
        avg = ave_list(lt)
        listSum = 0.0
        for i in lt:
            listSum += (i -avg) * (i - avg)
        return math.sqrt(listSum/(len(lt)-1))
    
            
from MSTATE import MSTATE

class IONISTATE():
    '''
    an ionization state of a protein
    '''


    def __init__(self, signs = [], residues = []):
        '''
        constructor
        '''
        
        # item numbers of signs and residues must be the same
        self.signs = signs
        self.residues = residues
        self.m_states = []
        self.partialRes = []            # contains part of the conformers, forming an ionization state
        self.ener       = []
        self.aveEner    = 0.0
        self.stdEner    = 0.0
        self.boltzAvg   = 0.0
    

    
    def get_micro_states(self):
        '''
        get all the microstates of this ionization state
        '''
        
        #print len(self.residues)
        
        # partialRes has the conformers with this ionization state
        for i in range(len(self.residues)):
            buff_res = []
            for conf in (self.residues[i]).conformers:
                if conf.is_this_sign(self.signs[i]):
                    buff_res.append(conf)
            #print len(buff_res)
            self.partialRes.append(buff_res)
            
        for confs in combine_list(self.partialRes):
            self.m_states.append(MSTATE(confs))    

                    
    def state_ener(self, fixProt):
        '''
        get the energy of this ionization state with a fix protein 'prot'
        '''
        self.get_micro_states()
        for state in self.m_states:
            state.get_total_ener(fixProt)
            self.ener.append(state.total_ener)

        
        self.aveEner = ave_list(self.ener)
        self.stdEner = std_list(self.ener)
        self.boltzAvg = boltz_ave_list(self.ener)
        
    def display_state(self):
        '''
        print the info of this state
        '''
        
        outStr = ''
        for sg in self.signs:
            outStr += ('%-10s' % sg)
            
        outStr += ('%.3f\t\t%.3f\t%.3f\t%d' % (self.boltzAvg, self.aveEner, self.stdEner, len(self.m_states)))
        
        return outStr 
                      
        
        
        
    
